/*
 * One example for NOI CSP-J Lesson 6:
 * <https://courses.fmsoft.cn/plzs/noijunior-search-algorithms.html>
 *
 * Author: Vincent Wei
 *  - <https://github.com/VincentWei>
 *  - <https://gitee.com/vincentwei7>
 *
 * Copyright (C) 2025 FMSoft <https://www.fmsoft.cn>.
 * License: GPLv3
 */
#include <stack>        // for stack
#include <queue>        // for queue

template <typename T>
class bin_tree_node {
  public:
    T payload;                          // 节点负载
    bin_tree_node *left;                // 左子节点
    bin_tree_node *right;               // 右子节点

    // 节点的构造函数
    bin_tree_node(const T& value) {
        payload = value;
        left = nullptr;
        right = nullptr;
    }

    // 节点的析构函数
    ~bin_tree_node() {
        if (left)
            delete left;
        if (right)
            delete right;
    }

    // 深度优先（depth-first）前序遍历
    template <typename context, typename visitor_func>
    void dfs_preorder(context* ctxt, visitor_func visitor) const
    {
        // call the visitor for the current node
        visitor(ctxt, payload);

        if (left)
            left->dfs_preorder(ctxt, visitor);
        if (right)
            right->dfs_preorder(ctxt, visitor);
    }

    // 深度优先（depth-first）中序遍历
    template <typename context, typename visitor_func>
    void dfs_inorder(context* ctxt, visitor_func visitor) const
    {
        if (left)
            left->dfs_inorder(ctxt, visitor);

        // call the visitor for the current node
        visitor(ctxt, payload);

        if (right)
            right->dfs_inorder(ctxt, visitor);
    }

    // 深度优先（depth-first）后序遍历
    template <typename context, typename visitor_func>
    void dfs_postorder(context* ctxt, visitor_func visitor) const
    {
        if (left)
            left->dfs_postorder(ctxt, visitor);
        if (right)
            right->dfs_postorder(ctxt, visitor);

        // call the visitor for the current node
        visitor(ctxt, payload);
    }

    // 广度优先（breadth-first）级序遍历
    template <typename context, typename visitor_func>
    void bfs(context* ctxt, visitor_func visitor) const
    {
        std::queue<const bin_tree_node*> queue;
        queue.push(this);

        while (!queue.empty()) {
            const bin_tree_node* node = queue.front();
            queue.pop();

            // call the visitor for the current node
            visitor(ctxt, node->payload);

            if (node->left)
                queue.push(node->left);
            if (node->right)
                queue.push(node->right);
        }
    }

    // 获取左子节点（const 版本）
    const bin_tree_node* left_child() const
    {
        return left;
    }

    // 获取左子节点（非 const 版本）
    bin_tree_node* left_child()
    {
        return left;
    }

    // 设置当前节点的左子节点，返回新节点；已有左子节点时返回 nullptr
    bin_tree_node* left_child(const T& value)
    {
        if (left)
            return nullptr;

        // Create a new node with the given value
        left = new bin_tree_node(value);
        return left;
    }

    // 移除当前节点的左子树
    bool erase_left_subtree()
    {
        if (left == nullptr)
            return false;

        delete left;
        left = nullptr;
        return true;
    }

    // 获取右子节点（const 版本）
    const bin_tree_node* right_child() const
    {
        return right;
    }

    // 获取右子节点（非 const 版本）
    bin_tree_node* right_child()
    {
        return right;
    }

    // 设置当前节点的右子节点，返回新节点；已有右子节点时返回 nullptr
    bin_tree_node* right_child(const T& value)
    {
        if (right)
            return nullptr;

        // Create a new node with the given value
        right = new bin_tree_node(value);
        return right;
    }

    // 移除当前节点的右子树
    bool erase_right_subtree()
    {
        if (right == nullptr)
            return false;

        delete right;
        right = nullptr;
        return true;
    }

    // 广度优先（breadth-first）层序遍历查找
    const bin_tree_node* bfs_find(const T& value) const
    {
        std::queue<const bin_tree_node*> queue;
        queue.push(this);

        while (!queue.empty()) {
            const bin_tree_node* node = queue.front();
            queue.pop();

            if (node->payload == value)
                return node;

            if (node->left)
                queue.push(node->left);
            if (node->right)
                queue.push(node->right);
        }

        return nullptr;
    }
};

template <class bin_tree_node, typename T>
bin_tree_node* insert(bin_tree_node* root, const T& value)
{
    bin_tree_node* temp = new bin_tree_node(value);

    // 如果 root 为空，则新节点为根节点
    if (root == nullptr)
        return temp;

    // 依据二叉搜索树的特点寻找新节点的父节点。
    bin_tree_node *parent = nullptr, *curr = root;
    while (curr != nullptr) {
        parent = curr;
        if (curr->payload > value)
            curr = curr->left;
        else if (curr->payload < value)
            curr = curr->right;
        else {
            delete temp;
            return root;        // 如果树中已有相同的负载值的节点
        }
    }

    // 如果 value 更小，则将 temp 作为父节点的左子节点，否则为右子节点。
    if (parent->payload > value)
        parent->left = temp;
    else
        parent->right = temp;

    return root;
}

template <class bin_tree_node, typename T>
bin_tree_node* search(bin_tree_node* root, const T& key)
{
    if (root == nullptr || root->payload == key)
        return root;

    // key 比 root 节点的负载值大
    if (root->payload < key)
        return search(root->right, key);

    // key 比 root 节点的负载值小
    return search(root->left, key);
}

// 注意：该函数并不是一个通用的中序后继节点获取函数。
// 该函数仅当右子节点不为空时工作，而这刚好是 BST 删除操作所满足的情形。
template <class bin_tree_node>
bin_tree_node* get_inorder_successor(bin_tree_node* curr)
{
    curr = curr->right;
    while (curr != nullptr && curr->left != nullptr)
        curr = curr->left;
    return curr;
}

template <class bin_tree_node, typename T>
bin_tree_node* delete_node(bin_tree_node* root, const T& key)
{
    // Base case
    if (root == nullptr)
        return root;

    // If key to be searched is in a subtree
    if (root->payload > key)
        root->left = delete_node(root->left, key);
    else if (root->payload < key)
        root->right = delete_node(root->right, key);
    // If root matches with the given key
    else {

        // 当 root 没有子节点或仅有右子节点时
        if (root->left == nullptr) {
            bin_tree_node* temp = root->right;
            delete root;
            return temp;
        }

        // 当 root 仅有左子节点时
        if (root->right == nullptr) {
            bin_tree_node* temp = root->left;
            delete root;
            return temp;
        }

        // 当 root 有两个子节点时
        bin_tree_node* succ = get_inorder_successor(root);
        root->payload = succ->payload;
        root->right = delete_node(root->right, succ->payload);
    }

    return root;
}

#include <iostream>     // for cin and cout
#include <sstream>      // for ostringstream
#include <string>       // for stod()
#include <cassert>      // for assert()

using namespace std;

using my_bin_tree_node = bin_tree_node<int>;
static void test_bst_search(my_bin_tree_node* root)
{
    for (int key = 5; key < 100; key += 15) {
        cout << "Key " << key;
        if (search(root, key)) {
            cout << " found\n";
        }
        else {
            cout << " not found\n";
        }
    }
}

int main()
{
    /* 创建如下的 BST
          50
         /  \
        30   70
       / \   / \
      20 40 60 80
    */

    my_bin_tree_node* root = new my_bin_tree_node(50);
    root = insert(root, 30);
    root = insert(root, 20);
    root = insert(root, 40);
    root = insert(root, 70);
    root = insert(root, 60);
    root = insert(root, 80);

    for (int key = 5; key < 100; key += 15) {
        cout << "Key " << key;
        if (root->bfs_find(key)) {
            cout << " found\n";
        }
        else {
            cout << " not found\n";
        }
    }
    struct context_print {
        bool root;
        ostream& os;
    };

    struct visitor_print {
        void operator() (context_print* ctxt, const double& value) {
            if (ctxt->root) {
                ctxt->os << value;
                ctxt->root = false;
            }
            else {
                ctxt->os << ", " << value;
            }
        }
    };

    context_print ctxt = { true, cout };
    clog << "Values in the tree:\n";
    ctxt.root = true;
    root->dfs_inorder(&ctxt, visitor_print{});
    cout << endl;

    test_bst_search(root);

    root = delete_node(root, 20);
    root = delete_node(root, 70);
    root = delete_node(root, 50);

    clog << "Values in the tree after deleting three nodes:\n";
    ctxt.root = true;
    root->dfs_inorder(&ctxt, visitor_print{});
    cout << endl;

    test_bst_search(root);

    delete root;
}

